#include <mm/heap.h>
#include <ds/semaphore.h>
#include <log.h>
#include <common.h>

#define CEIL(x, times) (((x) + (times - 1)) & ~(times - 1))

extern uptr __kernel_end;

void heap_init(heap_t* heap) {
    sem_init(&heap->sem, 1);

    header_t* head = heap->start;
    head->magic = MAGICNUMBER;
    head->used = 0;
    head->size = heap->end - heap->start - (sizeof(header_t) + sizeof(footer_t));

    footer_t* foot = (footer_t*)((uptr)heap->end - sizeof(footer_t));
    foot->magic = MAGICNUMBER;
    foot->header = head;
}

// expand the heap
void heap_expand(heap_t* heap, u32 size) {
    footer_t* old_foot = (footer_t*)((uptr)heap->end - sizeof(footer_t));
    heap->end = (void*)((uptr)heap->end + size);
    footer_t* new_foot = (footer_t*)((uptr)heap->end - sizeof(footer_t));

    header_t* head = old_foot->header;
    assert_msg(head->magic == MAGICNUMBER, "header magic mismatched when expanding");
    new_foot->magic = MAGICNUMBER;
    if (head->used == 0) {
        head->size = head->size + size;
        new_foot->header = head;
        old_foot->magic = 0;
        old_foot->header = 0;
        return;
    }

    head = (header_t*)((uptr)old_foot + sizeof(footer_t));
    head->magic = MAGICNUMBER;
    head->used = 0;
    head->size = size;
    new_foot->header = head;
}

void* alloc(heap_t* curr_heap, size_t size) {
    if (size >= (size_t)(curr_heap->end - curr_heap->start)) {
        panic("alloc() requested size is out of range!\n");
        return NULL;
    }
    sem_wait(&curr_heap->sem);
    uptr* ptr = (uptr*)curr_heap->start;
    void* ret = NULL;
    size = CEIL(size + sizeof(header_t) + sizeof(footer_t), 4);
    // | ---------------------------------heap------------------------------------- |
    // | head | ----------------------- head->size --------------------- | old_foot |
    // | head | -------------- size ---------------- | ----------------- | old_foot |
    // | head | ---- size ---- | new_foot | new_head | ----------------- | old_foot |
    header_t* head = (header_t*)ptr;
    for (;head < (header_t*)curr_heap->end;head = (header_t*)((uptr)head + head->size + sizeof(header_t) + sizeof(footer_t))) {
        if (head->used != 0 || head->size < size) continue;

        head->used = 1;
        footer_t* old_foot = (footer_t*)((uptr)head + sizeof(header_t) + head->size);
        footer_t* new_foot = (footer_t*)((uptr)head + size - sizeof(footer_t));
        header_t* new_head = (header_t*)((uptr)head + size);

        new_foot->header = head;
        new_foot->magic = MAGICNUMBER;

        old_foot->header = new_head;

        new_head->magic = MAGICNUMBER;
        new_head->size = head->size - size;
        new_head->used = 0;

        head->size = size - sizeof(footer_t) - sizeof(header_t);

        ret = (void*)((uptr)head + sizeof(header_t));
        break;
    }

    sem_post(&curr_heap->sem);

    if (ret == NULL) panic("alloc() no more room in heap\n");

    return ret;
}

void free(heap_t* curr_heap, void* ptr) {
    header_t* head = (header_t*)(ptr - sizeof(header_t));
    footer_t* footer = (footer_t*)(ptr + head->size);
    if (head->magic != MAGICNUMBER || footer->magic != MAGICNUMBER) {
        panic("free() invalid pointer\n");
        return;
    }
    footer_t* left_footer = (footer_t*)((uptr)head - sizeof(footer_t));

    header_t* right_header = (header_t*)((uptr)footer + sizeof(footer_t));
    footer_t* right_footer = (footer_t*)((uptr)right_header + right_header->size + sizeof(header_t));

    sem_wait(&curr_heap->sem);

    // merge into the right block if unused
    if ((void*)right_header < curr_heap->end &&
        (void*)right_header > curr_heap->start &&
        right_header->magic == MAGICNUMBER &&
        right_header->used == 0) {

        head->size += right_header->size + sizeof(header_t) + sizeof(footer_t);
        right_footer->header = head;

        footer->magic = 0;
        footer->header = 0;
        footer = right_footer;

        right_header->magic = 0;
        right_header->size = 0;
    }

    // merge into the left block if unused
    if ((void*)left_footer > curr_heap->start &&
        (void*)left_footer < curr_heap->end &&
        left_footer->magic == MAGICNUMBER &&
        left_footer->header->magic == MAGICNUMBER &&
        left_footer->header->used == 0) {

        left_footer->header->size += head->size + sizeof(header_t) + sizeof(footer_t);

        footer->header = left_footer->header;

        left_footer->magic = 0;
        left_footer->header = 0;
        head->magic = 0;
        head->size = 0;
    }

    head->used = 0;
    sem_post(&curr_heap->sem);
    return;
}